#include <disk.h>
#include <logf.h>
#include <string.h>
#include <types.h>
#include <kernel.h>

static disk_t disk_buffer[DISK_SIZE];

static void disk_send_cmd(disk_t *disk, uint32_t start_sector, uint32_t sector_count, int cmd)
{
    outb(DISK_DRIVE(disk), DISK_DRIVE_BASE | disk->type);
    outb(DISK_SECTOR(disk), (uint8_t)(sector_count >> 8));
    outb(DISK_LBA_L(disk), (uint8_t)(start_sector >> 24));
    outb(DISK_LBA_M(disk), 0);
    outb(DISK_LBA_H(disk), 0);

    outb(DISK_SECTOR(disk), (uint8_t)(sector_count >> 8));
    outb(DISK_LBA_L(disk), (uint8_t)start_sector);
    outb(DISK_LBA_M(disk), (uint8_t)(start_sector >> 8));
    outb(DISK_LBA_H(disk), (uint8_t)(start_sector >> 16));

    outb(DISK_CMD(disk), cmd);
}

static int disk_wait_data(disk_t *disk)
{
    uint8_t status;
    do {
        status = inb(DISK_STATUS(disk));
        if ((status & (DISK_STATUS_BUSY | DISK_STATUS_ERR | DISK_STATUS_DRQ)) != DISK_STATUS_BUSY) {
            break;
        }
    } while(TRUE);

    return (status & DISK_STATUS_ERR) ? -1 : 0;
}

static void disk_read_data(disk_t *disk, void *buffer, int size)
{
    uint16_t *c = (uint16_t *)buffer;
    for (int i = 0; i < size / 2; i++)
    {
        *c++ = inw(DISK_DATA(disk));
    }
}

static void disk_write_data(disk_t *disk, void *buffer, int size)
{
    uint16_t *c = (uint16_t *)buffer;
    for (int i = 0; i < size / 2; i++)
    {
        outw(DISK_DATA(disk), *c++);
    }
}

void init_disk()
{
    disk_t *disk = NULL;
    
    logf("checking disk...");
    kernel_memset(disk_buffer, 0, sizeof(disk_buffer));
    for (int i = 0; i < DISK_SIZE; i++)
    {
        disk = disk_buffer + i;
        kernel_sprintf(disk->name, "sd%c", i + 'a');
        disk->type = i == 0 ? DISK_MASTER : DISK_SLAVE;
        disk->port_base = IOBASE_PRIMARY;

        int err = identify_disk(disk);
        if (err == 0) 
            print_disk_info(disk);
    }
}

void read_part_info(disk_t *disk)
{
    disk_mbr_t mbr;
    disk_send_cmd(disk, 0, 1, DISK_CMD_READ);
    int err = disk_wait_data(disk);
    if (err < 0) {
        logf("read mbr info failed");
    }

    disk_read_data(disk, &mbr, sizeof(mbr));

    disk_part_info_t *pit = mbr.parts;
    disk_part_t *part = disk->parts + 1;
    for (int i = 0; i < DISK_PART_SIZE; i++, part++)
    {
        part->type = pit->type;
        if (part->type == FS_INVALID) {
            part->total_sector = 0;
            part->start_sector = 0;
            part->disk = NULL;
        } else {
            kernel_sprintf(part->name, "%s%d", disk->name, i);
            part->start_sector = pit->sector;
            part->total_sector = pit->total_sectors;
            part->disk = disk;
        }
    }
}

int identify_disk(disk_t *disk)
{
    disk_send_cmd(disk, 0, 0, DISK_CMD_IDENTIFY);
    int err = inb(DISK_STATUS(disk));
    if (err == 0) {
        logf("%s doesn't exists", disk->name);
        return -1;
    }
    
    err = disk_wait_data(disk);
    if (err < 0) {
        logf("disk[%s] read failed", disk->name);
        return -1;
    }

    uint16_t buf[256];
    disk_read_data(disk, buf, sizeof(buf));
    disk->sector_count = *(uint32_t *)(buf + 100);
    disk->sector_size = SECTOR_SIZE;

    disk_part_t *part = disk->parts + 0;
    part->disk = disk;
    kernel_sprintf(part->name, "%s%d", disk->name, 0);
    part->start_sector = 0;
    part->total_sector = disk->sector_count;
    part->type = FS_INVALID;

    // 读取分区表
    read_part_info(disk);
    return 0;
}

void print_disk_info(disk_t *disk)
{
    logf("disk name : %s", disk->name);
    logf("\tport base : %x", disk->port_base);
    logf("\ttotal size : %dM", disk->sector_count * disk->sector_size / 1024 / 1024);

    disk_part_t *part;
    for (int i = 0; i < DISK_PART_SIZE; i++)
    {
        part = disk->parts + i;
        if (part->type != FS_INVALID) {
            logf("\tname:%s\n\ttype: %x\n\tstart_sector: %d\n\ttotal_sector: %d\n",
                part->name, part->type, part->start_sector, part->total_sector);
        }
    }
}